# coding:utf-8
import torch


# 一、张量的拼接
# 1. torch.cat：不会产生新的维度,在原有维度上，将张量按dim维度进行拼接
# flag = True
flag = False
if flag:
    t = torch.ones((2, 3))

    t_0 = torch.cat([t, t], dim=0)  # 在行上进行扩展
    t_1 = torch.cat([t, t], dim=1)  # 在列上进行扩展
    print("t_0:{} shape:{}\nt_1:{} shape:{}".format(t_0, t_0.shape, t_1, t_1.shape))

# 2. torch.stack：会产生新的维度,在新的维度上，将张量在新的维度上进行扩展
# (第一维，第二维，第三维)
# 第一维：dim=0
flag = True
# flag = False
if flag:
    t = torch.ones((3, 4))
    # 新的维度为dim=0,占据第一维,在第一维上进行扩展
    t_stack = torch.stack([t, t], dim=0)  # 第一维扩展2个
    print("\nt_stack:{} shape:{}".format(t_stack, t_stack.shape))  # shape：(2,3,4)
    # 新的维度为dim=1,占据第二维，在第二维上进行扩展
    t_stack = torch.stack([t, t], dim=1)  # 在第二维扩展2个
    print("\nt_stack:{} shape:{}".format(t_stack, t_stack.shape))  # shape:(3, 2, 4)

    t_stack = torch.stack([t, t, t], dim=2)  # 在第三维扩展3个
    print("\nt_stack:{} shape:{}".format(t_stack, t_stack.shape))  # shape:(3,4,3)


# 二、张量的切分
# 1. torch.chunk：平均切分
# flag = True
flag = False
if flag:
    a = torch.ones(2, 5)
    # dim是切分维度，chunks是切分的份数
    list_of_tensors = torch.chunk(a, dim=1, chunks=2)
    for idx, t in enumerate(list_of_tensors):
        print("第{}个张量:{}，shape is {}".format(idx+1, t, t.shape))

# 2. torch.split：指定每一份的长度,每一份的长度可以不同
# flag = True
flag = False
if flag:
    a = torch.ones((2, 5))
    # list_of_tensors = torch.split(a, 2, dim=1)  # 每一份长度是int类型，int=2
    # for idx, t in enumerate(list_of_tensors):
    #     print("第{}个张量：{}, shape is {}".format(idx+1, t, t.shape))

    # 用list指定每一份的长度
    list_of_tensors = torch.split(a, [2, 1, 2], dim=1)
    for idx, t in enumerate((list_of_tensors)):
        print("第{}个张量：{}, shape is{}".format(idx+1, t, t.shape))


# 三、张量的索引
# 1. torch.index_select：根据索引号获取原张量的数据
flag = True
# flag = False
if flag:
    t = torch.randint(0, 9, size=(3, 3))  # 使用[0,9)的均匀分布来创建(3,3)的张量
    # idx = torch.tensor([0, 2], dtype=torch.long)  # 创建索引
    # idx = torch.tensor([0, 2], dtype=torch.float)  # 编译到index_select时报错，索引不能是float
    idx = torch.tensor([0, 2], dtype=torch.int64)  # 标准姿势
    t_select = torch.index_select(t, dim=0, index=idx)  # 在dim=0行维度上索引，每个索引取该维度上的索引号对应的所有数据
    print("t:\n{}\nt_select:\n{}".format(t, t_select))

# 2. torch.masked_select：根据bool_index来获取数据,返回一维张量
# flag = True
flag = False
if flag:
    t = torch.randint(0, 9, size=(3, 3))
    mask = t.ge(5)  # 生成掩码,布尔索引, >=5的为True
    t_select = torch.masked_select(t, mask)  # 取出mask为True对应位置的数据
    print("t:\n{}\nmask:\n{}t_select:\n{}".format(t, mask, t_select))

# 四、张量的形状变换
# 1. torch.reshape
# flag = True
flag = False
if flag:
    t = torch.randperm(8)  # 将[0, 8-1]里面的整数随机排列，返回一个一维张量
    t_reshape = torch.reshape(t, (2, 4))
    print("t:\n{}\nt_reshape:\n{}".format(t, t_reshape))

    print("*"*20)
    t_reshape = torch.reshape(t, (2, 2, -1))  # -1时，表示该维度的尺寸可以通过总数据量和其他维度数据量推倒出
    print("t:\n{}\nt_reshape:\n{}\nshape:\n{}".format(t, t_reshape, t_reshape.shape))

    print("-"*20)
    t[0] = 1024
    t_reshape = torch.reshape(t, (2, 4))
    print("t:\n", t, "\nt_reshape:\n", t_reshape)
    # 原数据与变换后的数据共享内存
    print(id(t.data), "\n", id(t_reshape.data))  # 注：是数据共享内存

# 2. torch.transpose：交换张量的任意两个维度
# flag = True
flag = False
if flag:
    t = torch.rand((2, 3, 4))  # 用[0,1)分布里面的随机数创建张量
    t_transpose = torch.transpose(t, dim0=1, dim1=2)  # 交换dim=1(第二维),dim=2(第三维)
    print("t shape:{}\nt_transpose shape:{}".format(t.shape, t_transpose.shape))

# 3. torch.squeeze：压缩长度为1的轴
# dim：若为None，移除所有长度为1的轴;若指定维度，当且仅当该轴长度为1时，可以被移除
# flag = True
flag = False
if flag:
    t = torch.rand((1, 2, 3, 1))
    t_sq = torch.squeeze(t)  # dim=None,默认移除所有长度为1的轴(也就是shape里面为1的维度)
    t_0 = torch.squeeze(t, dim=0)
    t_1 = torch.squeeze(t, dim=1)
    print(t.shape)
    # print(t_sq.shape)
    # print(t_0.shape)
    print(t_1.shape)


# 五、张量的数学运算
# (一)、加法
# 1. torch.add：
"""
功能：先乘后加,逐元素计算 input + alpha*other
input：第一个张量
alpha：乘项因子
other：第二个张量
"""
# flag = True
flag = False
if flag:
    t_0 = torch.randn((3, 3))
    t_1 = torch.ones_like(t_0)
    t_add = torch.add(t_0, 10, t_1)

    print("t_0:\n{}\nt_1:\n{}\nt_add_10:\n{}".format(t_0, t_1, t_add))